package com.hyj.langchain4j_springboot.control;

import com.hyj.langchain4j_springboot.config.AiConfig;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.service.TokenStream;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.LocalDateTime;

@RestController
@RequestMapping("/ai")
public class ChatController {
    @Autowired
    private QwenChatModel qwenChatModel;

    @Autowired
    private QwenStreamingChatModel qwenStreamingChatModel;

    @Autowired
    private AiConfig.Assistant assistant;

    @Autowired
    private AiConfig.UniqueAssistant uniqueAssistant;

    @RequestMapping("/chat")
    public String test(@RequestParam(defaultValue = "你是谁？")String input){
        return qwenChatModel.chat(input);
    }

    @RequestMapping(value = "/stream",produces = "text/stream;charset=UTF-8")
    public Flux<String> stream(@RequestParam(defaultValue = "你是谁？")String input){
        return Flux.create(fluxSink -> {
            qwenStreamingChatModel.chat(input, new StreamingChatResponseHandler() {
                @Override
                public void onPartialResponse(String s) {
                    fluxSink.next(s);
                }

                @Override
                public void onCompleteResponse(ChatResponse chatResponse) {
                    fluxSink.complete();
                }

                @Override
                public void onError(Throwable throwable) {
                    fluxSink.error(throwable);
                }
            });
        });
    }

    @RequestMapping("/chatM")
    public String chatWithMemory(@RequestParam(defaultValue = "你是谁？")String input){
        return assistant.chat(input);
    }

    @RequestMapping(value = "/streamM",produces = "text/stream;charset=UTF-8")
    public Flux<String> streamWithMemory(@RequestParam(defaultValue = "你是谁？")String input){
        TokenStream stream = assistant.stream(input);
//        TokenStream stream = assistant.stream(input, LocalDateTime.now().toString());
        return Flux.create(sink->{
            stream.onPartialResponse(sink::next)
                    .onCompleteResponse(c-> sink.complete())
                    .onError(sink::error)
                    .start();
        });
    }

    @RequestMapping("/chatMById")
    public String chatWithMemoryById(@RequestParam(defaultValue = "你是谁？")String input,@RequestParam Integer id){
        return uniqueAssistant.chat(id,input);
    }

    @RequestMapping(value = "/streamMById",produces = "text/stream;charset=UTF-8")
    public Flux<String> streamWithMemoryById(@RequestParam(defaultValue = "你是谁？")String input,@RequestParam Integer id){
        TokenStream stream = uniqueAssistant.stream(id,input);
        return Flux.create(sink->{
            stream.onPartialResponse(sink::next)
                    .onCompleteResponse(c-> sink.complete())
                    .onError(sink::error)
                    .start();
        });
    }
}
